% Iterative loop
nn = [200,100,50];
tt = [4,10,40,100,200,400];%
error_mat = zeros(size(tt,2),size(nn,2),1);


%Ns = paras.Ns;

for ttt = 1:size(tt,2)
    for n = 1:size(nn,2)
N = nn(n);T = tt(ttt);I = 3;
Ns = 25;
paras.T = T;paras.T = T;paras.N = N;paras.I = I;paras.Ns = Ns;

% Price Schedule
pvec = zeros(I,1,T);
pvec(1,1,:) = reshape(exp(linspace(0,2,T)),[1,1,T]);% (I,1,T)
pvec(2,1,:) = reshape(exp(linspace(0,1.7,T)),[1,1,T]);% (I,1,T)
pvec(3,1,:) = reshape(exp(linspace(0,1.4,T)),[1,1,T]);% (I,1,T)

% Generate artificial wage 

trend = (reshape(exp(linspace(0,log(10),T)),[1,1,T]));
rng(2,'twister')
wvec_tmp = lognrnd(3.58,0.7,1,N);
wvec = sort(wvec_tmp(:,:,1),2).*trend;
paras.uvec_init = wvec/10; % Initial value used to calculate the indirect utility of Comin.
paras.preference_type = 2;
paras.flag_sigma = 0;
paras.eta = -2;% from Auer et.al. (2021)
paras.eps_vec = [0.2;1;1.65];% Non-homothetic parametrization from Comin et.al (2021)
% Nh-CES parameters
if paras.flag_sigma == 0
    paras.sig = 0.25;% from Comin et.al (2021)
else
    paras.sig = 10;
end
medianw = median(wvec(1,:,1));
omega_tmp = 1./(medianw.^([0.2;1;1.65]*(1-paras.sig)));
omega_init = omega_tmp./sum(omega_tmp);
omega_init2 = omega_init(1:end-1,:);
options = optimoptions('fsolve', 'Algorithm', 'trust-region', 'TolFun', 1e-7,'Display','off');
omega_tmp2   =  fsolve( @(x) CalOmega(x,medianw,paras),omega_init2,options);  
paras.omega = [omega_tmp2;1-sum(omega_tmp2)];
if paras.preference_type ==1
    [v_vec, u_vec, b_vec]= AIDS(wvec,pvec,AIDSparas) ;
else
    [v_vec, u_vec, b_vec]= NhCES(wvec,pvec,paras) ;        
end


% Iterative
[U_vec] = CalMoneyMetric(wvec, b_vec, pvec,0);

ratio = reshape(U_vec./u_vec,[N,T]);
error_mat(ttt,n,1) = max(abs(log(ratio(:,T))));


    end
end
clearvars -except error_mat tt nn

save ../data/DataforMatlab/Fig2iterdata
%save ../data/DataforMatlab/Fig2iterdataFP


function [v_vec, u_vec, b_vec]= NhCES(wvec,pvec,paras) 
%%Nonhomothetic Preference: Given W and P, solve for u 
T = paras.T;
N = paras.N;
I = paras.I;
Ns = paras.Ns;
v_vec = zeros(1,N,T);
uvec_init = paras.uvec_init;
options = optimoptions('fsolve', 'TolFun', 1e-8, 'Display', 'off','Algorithm','Levenberg-Marquardt');

for t = 1:T
    for ns = 1: Ns   
        Index = (N/Ns)*(ns-1)+1:(N/Ns)*(ns);
       [ v_vec(:,Index,t),fval]  =  fsolve( @(x) NHU(x,wvec(:,Index,t),pvec(:,:,t),paras),uvec_init(:,Index),options);  % pid
    end 
end
%%Obtain EV
[u_vec,b_vec]= CalEV(v_vec,pvec,paras);

function err = NHU(v_vec,wvec,pvec,paras) 
if paras.flag_sigma ==1
    sig0 = paras.sig;
    paras.sig =sig0+log((v_vec).^paras.eta); % functional form from Auer (2021)
end
E = sum(paras.omega.*(((v_vec.^paras.eps_vec).*pvec).^(1-paras.sig))).^(1./(1-paras.sig));
err = abs(E-wvec);
end

function [u_vec,b_vec]= CalEV(v_vec,pvec,paras)
paras.sig =paras.sig;
sig0 = paras.sig;
if paras.flag_sigma ==1
     paras.sig =sig0+log((v_vec).^paras.eta);
end
pvec_b = pvec(:,:,1);
u_vec = sum(paras.omega.*(((v_vec.^paras.eps_vec).*pvec_b).^(1-paras.sig))).^(1./(1-paras.sig));
tmp = paras.omega.*(((v_vec.^paras.eps_vec).*pvec).^(1-paras.sig));
b_vec = tmp./sum(tmp);
end

end

function  err = CalOmega(omega_init,I_vec,paras) 
% function for solve for Omega 
options = optimoptions('fsolve', 'Algorithm', 'trust-region', 'TolFun', 1e-7,'Display','off');
 [ v_vec ]  =  fsolve( @(x) NHU(x,I_vec,omega_init,paras),I_vec/10,options);  
if paras.flag_sigma ==1
    sig0 = paras.sig;
    paras.sig =sig0+log((v_vec).^paras.eta);
end

omega = [omega_init;1-sum(omega_init)]; 
tmp = omega.*(((v_vec.^paras.eps_vec)).^(1-paras.sig));
B_vec = tmp./sum(tmp);

   err = (abs(B_vec(1:end,:)-1/paras.I));
    
function err = NHU(v_vec,I_vec,omega_init,paras) 
if paras.flag_sigma ==1
    sig0 = paras.sig;
    paras.sig =sig0+log((v_vec).^paras.eta);
end
omega = [omega_init;1-sum(omega_init)];
E = sum(omega.*(((v_vec.^paras.eps_vec)).^(1-paras.sig))).^(1./(1-paras.sig));
err = abs(E-I_vec);
end
end

